package com.analysis.client;

import com.analysis.config.BaiduAgentConfig;
import com.analysis.config.CommonConstant;
import com.analysis.dto.AiCallResponse;
import com.analysis.dto.Reference;
import com.analysis.exception.AppBuilderException;
import com.analysis.utils.LogUtils;
import com.analysis.utils.ParseTextUtils;
import com.baidubce.appbuilder.base.exception.AppBuilderServerException;
import com.baidubce.appbuilder.console.appbuilderclient.AppBuilderClient;
import com.baidubce.appbuilder.model.appbuilderclient.AppBuilderClientIterator;
import com.baidubce.appbuilder.model.appbuilderclient.AppBuilderClientResult;
import com.baidubce.appbuilder.model.appbuilderclient.Event;
import jakarta.annotation.PostConstruct;
import org.apache.commons.collections4.MapUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 百度千帆AppBuilder SDK客户端
 */
@Component
public class BaiduAgentAiChatClient {

    private static final Logger logger = LoggerFactory.getLogger(BaiduAgentAiChatClient.class);

    // 客户端缓存，避免重复创建
    private final ConcurrentHashMap<String, AppBuilderClient> clientCache = new ConcurrentHashMap<>();


    @Autowired
    private BaiduAgentConfig appBuilderConfig;

    @PostConstruct
    public void init() {
        // 设置系统属性（只设置一次）
        if (appBuilderConfig.getApiKey() != null) {
            System.setProperty("APPBUILDER_TOKEN", appBuilderConfig.getApiKey());
        }
    }

    /**
     * 获取或创建AppBuilderClient实例
     * 使用缓存机制提高效率
     */
    private AppBuilderClient getOrCreateClient(String appId) {
        return clientCache.computeIfAbsent(appId, id -> {
            if (appBuilderConfig.getEnableLogging()) {
                LogUtils.logInfo(logger, "创建新的AppBuilderClient实例", "appId", id);
            }
            return new AppBuilderClient(id);
        });
    }


    /**
     * 运行对话并获取完整答案（使用默认应用ID）
     * 优化版本：增加重试机制、超时控制、资源管理
     *
     * @param query 用户问题
     * @return 完整答案
     * @throws AppBuilderException 运行对话失败时抛出
     */
    public AiCallResponse callQuestion(String query) throws AppBuilderException {
        return call(query + CommonConstant.QUESTION_SIZE);
    }

    /**
     * 信息源分析
     */
    public List<Reference> callReference(String query) throws AppBuilderException {
        AiCallResponse contents = call("列出你为回答“" + query + "”所提供的内容所参考和引用的信息源, 返回信息源的title和url");
        logger.info("百度信源分析：" + contents);
        return ParseTextUtils.parseReferences(contents.getText());
    }

    /**
     * 运行对话并获取完整答案（指定应用ID）
     *
     * @param query 用户问题
     * @return 完整答案
     * @throws AppBuilderException 运行对话失败时抛出
     */
    public AiCallResponse call(String query) throws AppBuilderException {
        String appId = appBuilderConfig.getDefaultAppId();
        if (query == null || query.trim().isEmpty()) {
            throw new AppBuilderException("查询内容不能为空");
        }
        try {
            if (appBuilderConfig.getEnableLogging()) {
                LogUtils.logInfo(logger, "开始AppBuilder对话",
                        "appId", appId, "query", query.substring(0, Math.min(50, query.length())));
            }
            // 获取客户端实例
            AppBuilderClient builder = getOrCreateClient(appId);

            // 创建对话
            String conversationId = builder.createConversation();

            // 运行对话
            return runConversation(builder, query, conversationId);
        } catch (Exception e) {
            LogUtils.logError(logger, "AppBuilder对话失败", e);
            throw new AppBuilderException("对话失败: " + e.getMessage(), e);
        }
    }


    /**
     * 运行对话（带重试机制）
     */
    private AiCallResponse runConversation(AppBuilderClient builder, String query, String conversationId) throws IOException, AppBuilderServerException {
        StringBuilder answer = new StringBuilder();
        List<Reference> references = new ArrayList<>();
        double totalTokens = 0;

        // 执行
        AppBuilderClientIterator iterator = builder.run(query, conversationId, new String[]{}, false);

        while (iterator.hasNext()) {
            AppBuilderClientResult response = iterator.next();
            if (response != null && response.getAnswer() != null) {
                answer.append(response.getAnswer());
            }
            if (response != null && response.getEvents() != null) {
                // 解析events中的references
                parseReferencesFromEvents(response.getEvents(), references);

                // 解析events中的token用量
                totalTokens = totalTokens + extractTotalTokens(response.getEvents());
            }
        }

        String result = answer.toString();
        AiCallResponse aiCallResponse = new AiCallResponse();
        aiCallResponse.setText(result);
        aiCallResponse.setReferences(references);

        // 打印总token消耗
        if (totalTokens > 0) {
            System.err.println("baidu消耗token：" + totalTokens);
        }
        return aiCallResponse;
    }

    /**
     * 从events中解析references
     * 根据接口返回的数据格式，解析references的title和url
     *
     * @param events     事件数组
     * @param references 用于存储解析结果的references列表
     */
    private void parseReferencesFromEvents(Event[] events, List<Reference> references) {
        if (events == null || events.length == 0) {
            return;
        }

        for (Event event : events) {
            if (event == null || event.getDetail() == null) {
                continue;
            }

            // 检查是否是rag类型的事件，且包含references
            if ("rag".equals(event.getContentType())) {
                Map<String, Object> detail = event.getDetail();
                Object referencesObj = detail.get("references");
                if (referencesObj != null) {
                    @SuppressWarnings("unchecked")
                    List<Map<String, Object>> referencesList = (List<Map<String, Object>>) referencesObj;
                    for (Map<String, Object> refMap : referencesList) {
                        try {
                            Reference reference = new Reference();

                            // 获取title
                            Object titleObj = refMap.get("title");
                            if (titleObj != null) {
                                reference.setName(titleObj.toString());
                            }

                            // 获取url
                            Object urlObj = refMap.get("url");
                            if (urlObj != null) {
                                reference.setUrl(urlObj.toString());
                            }

                            // 只有当title和url都不为空时才添加到列表
                            if (reference.getName() != null && !reference.getName().trim().isEmpty() &&
                                    reference.getUrl() != null && !reference.getUrl().trim().isEmpty()) {
                                references.add(reference);

                                if (appBuilderConfig.getEnableLogging()) {
                                    LogUtils.logInfo(logger, "解析到reference",
                                            "title", reference.getName(),
                                            "url", reference.getUrl());
                                }
                            }
                        } catch (Exception e) {
                            if (appBuilderConfig.getEnableLogging()) {
                                LogUtils.logError(logger, "解析reference失败", "error", e.getMessage());
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * 仅从 events[].detail.usage.total_tokens 提取 token 总量。
     */
    public double extractTotalTokens(Event[] events) {
        double total = 0d;
        if (events != null) {
            for (Event event : events) {
                if (event == null || MapUtils.isEmpty(event.getUsage())) {
                    continue;
                }
                Map<String, Object> usage = event.getUsage();
                if (usage.containsKey("total_tokens")) {
                    try {
                        total = total + Double.parseDouble(String.valueOf(usage.get("total_tokens")));
                    } catch (Exception ignored) {
                    }
                }
            }
        }
        return total;
    }
}
